{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import transformers\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "from torch.utils.data import DataLoader, SequentialSampler\n",
    "\n",
    "\n",
    "import src.slurm\n",
    "import src.util\n",
    "from src.options import Options\n",
    "import src.data\n",
    "import src.evaluation\n",
    "import src.model\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "from loguru import logger\n",
    "# from src.model import FiDT5\n",
    "\n",
    "class FiDT5(src.model.FiDT5):\n",
    "    def generate(self, input_ids, attention_mask, max_length, \n",
    "                    num_beams=1, num_return_sequences=1, use_cache=False):\n",
    "        self.encoder.n_passages = input_ids.size(1)\n",
    "        \n",
    "        return transformers.T5ForConditionalGeneration.generate(\n",
    "            self,\n",
    "            input_ids=input_ids.view(input_ids.size(0), -1),\n",
    "            attention_mask=attention_mask.view(attention_mask.size(0), -1),\n",
    "            max_length=max_length,\n",
    "            num_beams=num_beams,\n",
    "            num_return_sequences=num_return_sequences,\n",
    "            use_cache=use_cache,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = transformers.T5Tokenizer.from_pretrained(\"/home/xionggm/codes/decode-answer-logical-form/PLMs/t5-base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "config.json\t\tmodel.safetensors  spiece.model\n",
      "generation_config.json\tpytorch_model.bin  tokenizer.json\n"
     ]
    }
   ],
   "source": [
    "!ls /home/xionggm/codes/decode-answer-logical-form/PLMs/t5-base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "t5 = transformers.T5ForConditionalGeneration.from_pretrained(\"/home/xionggm/codes/decode-answer-logical-form/PLMs/t5-base\", return_dict=False)\n",
    "model = FiDT5(t5.config)\n",
    "pt = torch.load(os.path.join(\"save-FiD/WebQSP-t5-base/checkpoint/best_dev\", \"model.pt\"), map_location='cpu')\n",
    "model.load_state_dict(pt)\n",
    "model = model.to(\"cpu\")\n",
    "del t5, pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = [\"Translate the following English text to French: 'Hello, how are you?'\",\"Translate the following English text to French: 'Hello, how old are you?'\"]\n",
    "# input_text = \"hello!!!!\"\n",
    "input_ids = tokenizer.encode(input_text, return_tensors=\"pt\")\n",
    "attention_mask = torch.ones(input_ids.size(), dtype=torch.long, device=input_ids.device)\n",
    "\n",
    "# Generate output\n",
    "n=2\n",
    "output_ids = model.generate(input_ids,num_beams=n,attention_mask=attention_mask,num_return_sequences=n,max_length=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['heim moi moi', 'heim moi']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "# Decode and print the generated text\n",
    "output_texts = [tokenizer.decode(output_id, skip_special_tokens=True, clean_up_tokenization_spaces=True) for output_id in output_ids]\n",
    "\n",
    "output_texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "t5 = transformers.T5ForConditionalGeneration.from_pretrained(\n",
    "            \"/home/xionggm/codes/decode-answer-logical-form/PLMs/t5-small\",\n",
    "            return_dict=False,\n",
    "        )\n",
    "model = FiDT5(t5.config)\n",
    "model.load_state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from transformers import T5ForConditionalGeneration, T5Tokenizer\n",
    "import torch\n",
    "model_name = \"PLMs/t5-base\"\n",
    "tokenizer = T5Tokenizer.from_pretrained(model_name)\n",
    "model = T5ForConditionalGeneration.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xionggm/miniconda3/lib/python3.11/site-packages/transformers/generation/utils.py:1355: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['', '', '', '', '', '', '', '', '', '']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_text = [\"Translate the following English text to French: 'Hello, how are you?'\",\"Translate the following English text to French: 'Hello, how old are you?'\"]\n",
    "input_ids = tokenizer.encode(input_text, return_tensors=\"pt\")\n",
    "\n",
    "# Generate output\n",
    "output_ids = model.generate(input_ids,num_beams=10,num_return_sequences=10)\n",
    "\n",
    "# Decode and print the generated text\n",
    "output_texts = [tokenizer.decode(output_id, skip_special_tokens=True, clean_up_tokenization_spaces=True) for output_id in output_ids]\n",
    "\n",
    "output_texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'list' object has no attribute 'shape'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[1;32m/home/xionggm/codes/decode-answer-logical-form/test.ipynb Cell 11\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> <a href='vscode-notebook-cell://ssh-remote%2B7b22686f73744e616d65223a22413130307832227d/home/xionggm/codes/decode-answer-logical-form/test.ipynb#X13sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a>\u001b[0m input_text\u001b[39m.\u001b[39;49mshape\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'list' object has no attribute 'shape'"
     ]
    }
   ],
   "source": [
    "input_text.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 20])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_ids.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([100, 200])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids = torch.LongTensor(torch.randint(0, 1000, (100, 200)))\n",
    "input_ids.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_ids = model.generate(input_ids,num_beams=10,num_return_sequences=10,max_new_tokens=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
